{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| eval: false\n",
    "! [ -e /content ] && pip install -Uqq fastai  # upgrade fastai on colab"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "skip_exec: true\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "from __future__ import annotations\n",
    "from fastai.basics import *\n",
    "from fastai.callback.progress import *\n",
    "from fastai.text.data import TensorText\n",
    "from fastai.tabular.all import TabularDataLoaders, Tabular\n",
    "from fastai.callback.hook import total_params\n",
    "from fastai.callback.tracker import SaveModelCallback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp callback.wandb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wandb\n",
    "\n",
    "> Integration with [Weights & Biases](https://docs.wandb.com/library/integrations/fastai) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First thing first, you need to install wandb with\n",
    "```\n",
    "pip install wandb\n",
    "```\n",
    "Create a free account then run \n",
    "``` \n",
    "wandb login\n",
    "```\n",
    "in your terminal. Follow the link to get an API token that you will need to paste, then you're all set!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class WandbCallback(Callback):\n",
    "    \"Saves model topology, losses & metrics\"\n",
    "    remove_on_fetch,order = True,Recorder.order+1\n",
    "    # Record if watch has been called previously (even in another instance)\n",
    "    _wandb_watch_called = False\n",
    "\n",
    "    def __init__(self, \n",
    "                 log:str=None, # What to log (can be `gradients`, `parameters`, `all` or None) \n",
    "                 log_preds:bool=True, # Whether to log model predictions on a `wandb.Table`\n",
    "                 log_preds_every_epoch:bool=False, # Whether to log predictions every epoch or at the end\n",
    "                 log_model:bool=False, # Whether to save the model checkpoint to a `wandb.Artifact` \n",
    "                 model_name:str=None, # The name of the `model_name` to save, overrides `SaveModelCallback`\n",
    "                 log_dataset:bool=False, # Whether to log the dataset to a `wandb.Artifact`\n",
    "                 dataset_name:str=None, # A name to log the dataset with\n",
    "                 valid_dl:TfmdDL=None, # If `log_preds=True`, then the samples will be drawn from `valid_dl`\n",
    "                 n_preds:int=36, # How many samples to log predictions \n",
    "                 seed:int=12345, # The seed of the samples drawn\n",
    "                 reorder=True):\n",
    "        store_attr()\n",
    "    \n",
    "    def after_create(self):\n",
    "        # log model\n",
    "        if self.log_model:\n",
    "            if not hasattr(self, 'save_model'):\n",
    "                # does not have the SaveModelCallback\n",
    "                self.learn.add_cb(SaveModelCallback(fname=ifnone(self.model_name, 'model')))\n",
    "            else:\n",
    "                # override SaveModelCallback\n",
    "                if self.model_name is not None:\n",
    "                    self.save_model.fname = self.model_name\n",
    "            \n",
    "    def before_fit(self):\n",
    "        \"Call watch method to log model topology, gradients & weights\"\n",
    "        # Check if wandb.init has been called\n",
    "        if wandb.run is None:\n",
    "            raise ValueError('You must call wandb.init() before WandbCallback()')\n",
    "        # W&B log step\n",
    "        self._wandb_step = wandb.run.step - 1  # -1 except if the run has previously logged data (incremented at each batch)\n",
    "        self._wandb_epoch = 0 if not(wandb.run.step) else math.ceil(wandb.run.summary['epoch']) # continue to next epoch\n",
    "        \n",
    "        self.run = not hasattr(self.learn, 'lr_finder') and not hasattr(self, \"gather_preds\") and rank_distrib()==0\n",
    "        if not self.run: return\n",
    "\n",
    "        # Log config parameters\n",
    "        log_config = self.learn.gather_args()\n",
    "        _format_config(log_config)\n",
    "        try:\n",
    "            wandb.config.update(log_config, allow_val_change=True)\n",
    "        except Exception as e:\n",
    "            print(f'WandbCallback could not log config parameters -> {e}')\n",
    "\n",
    "        if not WandbCallback._wandb_watch_called:\n",
    "            WandbCallback._wandb_watch_called = True\n",
    "            # Logs model topology and optionally gradients and weights\n",
    "            if self.log is not None:\n",
    "                wandb.watch(self.learn.model, log=self.log)\n",
    "\n",
    "        # log dataset\n",
    "        assert isinstance(self.log_dataset, (str, Path, bool)), 'log_dataset must be a path or a boolean'\n",
    "        if self.log_dataset is True:\n",
    "            if Path(self.dls.path) == Path('.'):\n",
    "                print('WandbCallback could not retrieve the dataset path, please provide it explicitly to \"log_dataset\"')\n",
    "                self.log_dataset = False\n",
    "            else:\n",
    "                self.log_dataset = self.dls.path\n",
    "        if self.log_dataset:\n",
    "            self.log_dataset = Path(self.log_dataset)\n",
    "            assert self.log_dataset.is_dir(), f'log_dataset must be a valid directory: {self.log_dataset}'\n",
    "            metadata = {'path relative to learner': os.path.relpath(self.log_dataset, self.learn.path)}\n",
    "            log_dataset(path=self.log_dataset, name=self.dataset_name, metadata=metadata)\n",
    "\n",
    "        if self.log_preds:\n",
    "            try:\n",
    "                if not self.valid_dl:\n",
    "                    #Initializes the batch watched\n",
    "                    wandbRandom = random.Random(self.seed)  # For repeatability\n",
    "                    self.n_preds = min(self.n_preds, len(self.dls.valid_ds))\n",
    "                    idxs = wandbRandom.sample(range(len(self.dls.valid_ds)), self.n_preds)\n",
    "                    if isinstance(self.dls,  TabularDataLoaders):\n",
    "                        test_items = getattr(self.dls.valid_ds.items, 'iloc', self.dls.valid_ds.items)[idxs]\n",
    "                        self.valid_dl = self.dls.test_dl(test_items, with_labels=True, process=False)\n",
    "                    else:\n",
    "                        test_items = [getattr(self.dls.valid_ds.items, 'iloc', self.dls.valid_ds.items)[i] for i in idxs]\n",
    "                        self.valid_dl = self.dls.test_dl(test_items, with_labels=True)\n",
    "                self.learn.add_cb(FetchPredsCallback(dl=self.valid_dl, with_input=True, with_decoded=True, reorder=self.reorder))\n",
    "            except Exception as e:\n",
    "                self.log_preds = False\n",
    "                print(f'WandbCallback was not able to prepare a DataLoader for logging prediction samples -> {e}')\n",
    "        \n",
    "    def before_batch(self): \n",
    "        self.ti_batch = time.perf_counter()\n",
    "        \n",
    "    def after_batch(self):\n",
    "        \"Log hyper-parameters and training loss\"\n",
    "        if self.training:\n",
    "            batch_time = time.perf_counter() - self.ti_batch\n",
    "            self._wandb_step += 1\n",
    "            self._wandb_epoch += 1/self.n_iter\n",
    "            hypers = {f'{k}_{i}':v for i,h in enumerate(self.opt.hypers) for k,v in h.items()}\n",
    "            wandb.log({'epoch': self._wandb_epoch, 'train_loss': self.smooth_loss, 'raw_loss': self.loss, **hypers}, step=self._wandb_step)\n",
    "            wandb.log({'train_samples_per_sec': len(self.xb[0]) / batch_time}, step=self._wandb_step)\n",
    "            \n",
    "    def log_predictions(self):\n",
    "        try:\n",
    "            inp,preds,targs,out = self.learn.fetch_preds.preds\n",
    "            b = tuplify(inp) + tuplify(targs)\n",
    "            x,y,its,outs = self.valid_dl.show_results(b, out, show=False, max_n=self.n_preds)\n",
    "            wandb.log(wandb_process(x, y, its, outs, preds), step=self._wandb_step)\n",
    "        except Exception as e:\n",
    "            self.log_preds = False\n",
    "            self.remove_cb(FetchPredsCallback)\n",
    "            print(f'WandbCallback was not able to get prediction samples -> {e}')\n",
    "    \n",
    "\n",
    "    def after_epoch(self):\n",
    "        \"Log validation loss and custom metrics & log prediction samples\"\n",
    "        # Correct any epoch rounding error and overwrite value\n",
    "        self._wandb_epoch = round(self._wandb_epoch)\n",
    "        if self.log_preds and self.log_preds_every_epoch:\n",
    "            self.log_predictions()\n",
    "        wandb.log({'epoch': self._wandb_epoch}, step=self._wandb_step)\n",
    "        wandb.log({n:s for n,s in zip(self.recorder.metric_names, self.recorder.log) if n not in ['train_loss', 'epoch', 'time']}, step=self._wandb_step)\n",
    "    \n",
    "    \n",
    "    \n",
    "    def after_fit(self):\n",
    "        if self.log_preds and not self.log_preds_every_epoch:\n",
    "            self.log_predictions()\n",
    "        if self.log_model:\n",
    "            if self.save_model.last_saved_path is None:\n",
    "                print('WandbCallback could not retrieve a model to upload')\n",
    "            else:\n",
    "                metadata = {n:s for n,s in zip(self.recorder.metric_names, self.recorder.log) if n not in ['train_loss', 'epoch', 'time']}\n",
    "                log_model(self.save_model.last_saved_path, name=self.save_model.fname, metadata=metadata)                \n",
    "        self.run = True\n",
    "        if self.log_preds: self.remove_cb(FetchPredsCallback)\n",
    "        \n",
    "        wandb.log({})  # ensure sync of last step\n",
    "        self._wandb_step += 1\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "### WandbCallback\n",
       "\n",
       ">      WandbCallback (log:'str'=None, log_preds:'bool'=True,\n",
       ">                     log_preds_every_epoch:'bool'=False,\n",
       ">                     log_model:'bool'=False, model_name:'str'=None,\n",
       ">                     log_dataset:'bool'=False, dataset_name:'str'=None,\n",
       ">                     valid_dl:'TfmdDL'=None, n_preds:'int'=36,\n",
       ">                     seed:'int'=12345, reorder=True)\n",
       "\n",
       "Saves model topology, losses & metrics\n",
       "\n",
       "|    | **Type** | **Default** | **Details** |\n",
       "| -- | -------- | ----------- | ----------- |\n",
       "| log | str | None | What to log (can be `gradients`, `parameters`, `all` or None) |\n",
       "| log_preds | bool | True | Whether to log model predictions on a `wandb.Table` |\n",
       "| log_preds_every_epoch | bool | False | Whether to log predictions every epoch or at the end |\n",
       "| log_model | bool | False | Whether to save the model checkpoint to a `wandb.Artifact` |\n",
       "| model_name | str | None | The name of the `model_name` to save, overrides `SaveModelCallback` |\n",
       "| log_dataset | bool | False | Whether to log the dataset to a `wandb.Artifact` |\n",
       "| dataset_name | str | None | A name to log the dataset with |\n",
       "| valid_dl | TfmdDL | None | If `log_preds=True`, then the samples will be drawn from `valid_dl` |\n",
       "| n_preds | int | 36 | How many samples to log predictions |\n",
       "| seed | int | 12345 | The seed of the samples drawn |\n",
       "| reorder | bool | True |  |"
      ],
      "text/plain": [
       "<nbdev.showdoc.BasicMarkdownRenderer>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(WandbCallback)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optionally logs weights and or gradients depending on `log` (can be \"gradients\", \"parameters\", \"all\" or None), sample predictions if ` log_preds=True` that will come from `valid_dl` or a random sample of the validation set (determined by `seed`). `n_preds` are logged in this case.\n",
    "\n",
    "If used in combination with `SaveModelCallback`, the best model is saved as well (can be deactivated with `log_model=False`).\n",
    "\n",
    "Datasets can also be tracked:\n",
    "\n",
    "* if `log_dataset` is `True`, tracked folder is retrieved from `learn.dls.path`\n",
    "* `log_dataset` can explicitly be set to the folder to track\n",
    "* the name of the dataset can explicitly be given through `dataset_name`, otherwise it is set to the folder name\n",
    "* *Note: the subfolder \"models\" is always ignored*\n",
    "\n",
    "For custom scenarios, you can also manually use functions `log_dataset` and `log_model` to respectively log your own datasets and models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch\n",
    "def gather_args(self:Learner):\n",
    "    \"Gather config parameters accessible to the learner\"\n",
    "    # args stored by `store_attr`\n",
    "    cb_args = {f'{cb}':getattr(cb,'__stored_args__',True) for cb in self.cbs}\n",
    "    args = {'Learner':self, **cb_args}\n",
    "    # input dimensions\n",
    "    try:\n",
    "        n_inp = self.dls.train.n_inp\n",
    "        args['n_inp'] = n_inp\n",
    "        xb = self.dls.valid.one_batch()[:n_inp]\n",
    "        args.update({f'input {n+1} dim {i+1}':d for n in range(n_inp) for i,d in enumerate(list(detuplify(xb[n]).shape))})\n",
    "    except: print(f'Could not gather input dimensions')\n",
    "    # other useful information\n",
    "    with ignore_exceptions():\n",
    "        args['batch size'] = self.dls.bs\n",
    "        args['batch per epoch'] = len(self.dls.train)\n",
    "        args['model parameters'] = total_params(self.model)[0]\n",
    "        args['device'] = self.dls.device.type\n",
    "        args['frozen'] = bool(self.opt.frozen_idx)\n",
    "        args['frozen idx'] = self.opt.frozen_idx\n",
    "        args['dataset.tfms'] = f'{self.dls.dataset.tfms}'\n",
    "        args['dls.after_item'] = f'{self.dls.after_item}'\n",
    "        args['dls.before_batch'] = f'{self.dls.before_batch}'\n",
    "        args['dls.after_batch'] = f'{self.dls.after_batch}'\n",
    "    return args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "#### Learner.gather_args\n",
       "\n",
       ">      Learner.gather_args ()\n",
       "\n",
       "Gather config parameters accessible to the learner"
      ],
      "text/plain": [
       "<nbdev.showdoc.BasicMarkdownRenderer>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(Learner.gather_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _make_plt(img):\n",
    "    \"Make plot to image resolution\"\n",
    "    # from https://stackoverflow.com/a/13714915\n",
    "    my_dpi = 100\n",
    "    fig = plt.figure(frameon=False, dpi=my_dpi)\n",
    "    h, w = img.shape[:2]\n",
    "    fig.set_size_inches(w / my_dpi, h / my_dpi)\n",
    "    ax = plt.Axes(fig, [0., 0., 1., 1.])\n",
    "    ax.set_axis_off()\n",
    "    fig.add_axes(ax)\n",
    "    return fig, ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _format_config_value(v):\n",
    "    if isinstance(v, list):\n",
    "        return [_format_config_value(item) for item in v]\n",
    "    elif hasattr(v, '__stored_args__'):\n",
    "        return {**_format_config(v.__stored_args__), '_name': v}\n",
    "    return v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _format_config(config):\n",
    "    \"Format config parameters before logging them\"\n",
    "    for k,v in config.items():\n",
    "        if isinstance(v, dict):\n",
    "            config[k] = _format_config(v)\n",
    "        else:\n",
    "            config[k] = _format_config_value(v)\n",
    "    return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _format_metadata(metadata):\n",
    "    \"Format metadata associated to artifacts\"\n",
    "    for k,v in metadata.items(): metadata[k] = str(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def log_dataset(path, name=None, metadata={}, description='raw dataset'):\n",
    "    \"Log dataset folder\"\n",
    "    # Check if wandb.init has been called in case datasets are logged manually\n",
    "    if wandb.run is None:\n",
    "        raise ValueError('You must call wandb.init() before log_dataset()')\n",
    "    path = Path(path)\n",
    "    if not path.is_dir():\n",
    "        raise f'path must be a valid directory: {path}'\n",
    "    name = ifnone(name, path.name)\n",
    "    _format_metadata(metadata)\n",
    "    artifact_dataset = wandb.Artifact(name=name, type='dataset', metadata=metadata, description=description)\n",
    "    # log everything except \"models\" folder\n",
    "    for p in path.ls():\n",
    "        if p.is_dir():\n",
    "            if p.name != 'models': artifact_dataset.add_dir(str(p.resolve()), name=p.name)\n",
    "        else: artifact_dataset.add_file(str(p.resolve()))\n",
    "    wandb.run.use_artifact(artifact_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "#### log_dataset\n",
       "\n",
       ">      log_dataset (path, name=None, metadata={}, description='rawdataset')\n",
       "\n",
       "Log dataset folder"
      ],
      "text/plain": [
       "<nbdev.showdoc.BasicMarkdownRenderer>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(log_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def log_model(path, name=None, metadata={}, description='trained model'):\n",
    "    \"Log model file\"\n",
    "    if wandb.run is None:\n",
    "        raise ValueError('You must call wandb.init() before log_model()')\n",
    "    path = Path(path)\n",
    "    if not path.is_file():\n",
    "        raise f'path must be a valid file: {path}'\n",
    "    name = ifnone(name, f'run-{wandb.run.id}-model')\n",
    "    _format_metadata(metadata)    \n",
    "    artifact_model = wandb.Artifact(name=name, type='model', metadata=metadata, description=description)\n",
    "    with artifact_model.new_file(str(Path(name).with_suffix(\".pth\")), mode='wb') as fa:\n",
    "        fa.write(path.read_bytes())\n",
    "    wandb.run.log_artifact(artifact_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "#### log_model\n",
       "\n",
       ">      log_model (path, name=None, metadata={}, description='trainedmodel')\n",
       "\n",
       "Log model file"
      ],
      "text/plain": [
       "<nbdev.showdoc.BasicMarkdownRenderer>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(log_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@dispatch\n",
    "def wandb_process(x:TensorImage, y, samples, outs, preds):\n",
    "    \"Process `sample` and `out` depending on the type of `x/y`\"\n",
    "    res_input, res_pred, res_label = [],[],[]\n",
    "    for s,o in zip(samples, outs):\n",
    "        img = s[0].permute(1,2,0)\n",
    "        res_input.append(wandb.Image(img, caption='Input_data'))\n",
    "        for t, capt, res in ((o[0], \"Prediction\", res_pred), (s[1], \"Ground_Truth\", res_label)):\n",
    "            fig, ax = _make_plt(img)\n",
    "            # Superimpose label or prediction to input image\n",
    "            ax = img.show(ctx=ax)\n",
    "            ax = t.show(ctx=ax)\n",
    "            res.append(wandb.Image(fig, caption=capt))\n",
    "            plt.close(fig)\n",
    "    return {\"Inputs\":res_input, \"Predictions\":res_pred, \"Ground_Truth\":res_label}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _unlist(l):\n",
    "    \"get element of lists of lenght 1\"\n",
    "    if isinstance(l, (list, tuple)):\n",
    "        if len(l) == 1: return l[0]\n",
    "    else: return l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@dispatch\n",
    "def wandb_process(x:TensorImage, y:TensorCategory|TensorMultiCategory, samples, outs, preds):\n",
    "    table = wandb.Table(columns=[\"Input image\", \"Ground_Truth\", \"Predictions\"])\n",
    "    for (image, label), pred_label in zip(samples,outs):\n",
    "        table.add_data(wandb.Image(image.permute(1,2,0)), label, _unlist(pred_label))\n",
    "    return {\"Prediction_Samples\": table}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@dispatch\n",
    "def wandb_process(x:TensorImage, y:TensorMask, samples, outs, preds):\n",
    "    res = []\n",
    "    codes = getattr(outs[0][0], 'codes', None)\n",
    "    if codes is not None:\n",
    "        class_labels = [{'name': name, 'id': id}  for id, name in enumerate(codes)] \n",
    "    else:\n",
    "        class_labels = [{'name': i, 'id': i} for i in range(preds.shape[1])]\n",
    "    table = wandb.Table(columns=[\"Input Image\", \"Ground_Truth\", \"Predictions\"])\n",
    "    for (image, label), pred_label in zip(samples, outs):\n",
    "        img = image.permute(1,2,0)\n",
    "        table.add_data(wandb.Image(img),\n",
    "                       wandb.Image(img, masks={\"Ground_Truth\": {'mask_data': label.numpy().astype(np.uint8)}}, classes=class_labels), \n",
    "                       wandb.Image(img, masks={\"Prediction\":   {'mask_data': pred_label[0].numpy().astype(np.uint8)}}, classes=class_labels) \n",
    "                      )\n",
    "    return {\"Prediction_Samples\": table}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@dispatch\n",
    "def wandb_process(x:TensorText, y:TensorCategory|TensorMultiCategory, samples, outs, preds):\n",
    "    data = [[s[0], s[1], o[0]] for s,o in zip(samples,outs)]\n",
    "    return {\"Prediction_Samples\": wandb.Table(data=data, columns=[\"Text\", \"Target\", \"Prediction\"])}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@dispatch\n",
    "def wandb_process(x:Tabular, y:Tabular, samples, outs, preds):\n",
    "    df = x.all_cols\n",
    "    for n in x.y_names: df[n+'_pred'] = y[n].values\n",
    "    return {\"Prediction_Samples\": wandb.Table(dataframe=df)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example of use:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once your have defined your `Learner`, before you call to `fit` or `fit_one_cycle`, you need to initialize wandb:\n",
    "```\n",
    "import wandb\n",
    "wandb.init()\n",
    "```\n",
    "To use Weights & Biases without an account, you can call `wandb.init(anonymous='allow')`.\n",
    "\n",
    "Then you add the callback to your `learner` or call to `fit` methods, potentially with `SaveModelCallback` if you want to save the best model:\n",
    "```\n",
    "from fastai.callback.wandb import *\n",
    "\n",
    "# To log only during one training phase\n",
    "learn.fit(..., cbs=WandbCallback())\n",
    "\n",
    "# To log continuously for all training phases\n",
    "learn = learner(..., cbs=WandbCallback())\n",
    "```\n",
    "Datasets and models can be tracked through the callback or directly through `log_model` and `log_dataset` functions.\n",
    "\n",
    "For more details, refer to [W&B documentation](https://docs.wandb.com/library/integrations/fastai)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| slow\n",
    "from fastai.vision.all import *\n",
    "import tempfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.12.17"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "W&B syncing is set to <code>`offline`<code> in this directory.  <br/>Run <code>`wandb online`<code> or set <code>WANDB_MODE=online<code> to enable cloud syncing."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.707184</td>\n",
       "      <td>0.196000</td>\n",
       "      <td>00:02</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.707796</td>\n",
       "      <td>0.522665</td>\n",
       "      <td>00:02</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.599319</td>\n",
       "      <td>0.323519</td>\n",
       "      <td>00:02</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Better model found at epoch 0 with valid_loss value: 0.32351893186569214.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.854327</td>\n",
       "      <td>0.279063</td>\n",
       "      <td>00:02</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Better model found at epoch 0 with valid_loss value: 0.279062956571579.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='0.000 MB of 0.000 MB uploaded (0.000 MB deduped)\\r'), FloatProgress(value=1.0, max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>epoch</td><td>▁▁▁▁▂▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███</td></tr><tr><td>eps_0</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>eps_1</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>eps_2</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>lr_0</td><td>██████████▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>lr_1</td><td>██████████▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>lr_2</td><td>▁▁▁▁▁▁▁▁▁▁██████████████████████████████</td></tr><tr><td>mom_0</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>mom_1</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>mom_2</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>raw_loss</td><td>█▄▃▃▁▁▂▄▂▄▇▄▃▂▄▁▃▃▃▂▆▄▄▂▃▃▁▁▂▁█▄▃▂▃▂▂▇▃▂</td></tr><tr><td>sqr_mom_0</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>sqr_mom_1</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>sqr_mom_2</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>train_loss</td><td>▇▅▄▃▃▂▂▂▂▂▇▅▄▃▃▂▂▂▂▂▅▄▄▃▃▂▂▁▁▁█▅▄▃▃▂▂▃▂▃</td></tr><tr><td>train_samples_per_sec</td><td>▃▇███▇▇███▃▄▇███▇▇▇█▃▇▇▅▆▆▆███▁▇█▇▇▇▆▆██</td></tr><tr><td>valid_loss</td><td>▁█▄▃</td></tr><tr><td>wd_0</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>wd_1</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>wd_2</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>epoch</td><td>4</td></tr><tr><td>eps_0</td><td>1e-05</td></tr><tr><td>eps_1</td><td>1e-05</td></tr><tr><td>eps_2</td><td>1e-05</td></tr><tr><td>lr_0</td><td>0.0005</td></tr><tr><td>lr_1</td><td>0.0005</td></tr><tr><td>lr_2</td><td>0.005</td></tr><tr><td>mom_0</td><td>0.9</td></tr><tr><td>mom_1</td><td>0.9</td></tr><tr><td>mom_2</td><td>0.9</td></tr><tr><td>raw_loss</td><td>0.47733</td></tr><tr><td>sqr_mom_0</td><td>0.99</td></tr><tr><td>sqr_mom_1</td><td>0.99</td></tr><tr><td>sqr_mom_2</td><td>0.99</td></tr><tr><td>train_loss</td><td>0.85433</td></tr><tr><td>train_samples_per_sec</td><td>3029.923</td></tr><tr><td>valid_loss</td><td>0.27906</td></tr><tr><td>wd_0</td><td>0.01</td></tr><tr><td>wd_1</td><td>0.01</td></tr><tr><td>wd_2</td><td>0.01</td></tr></table><br/></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "You can sync this run to the cloud by running:<br/><code>wandb sync /tmp/tmpcbvaf5w6/wandb/offline-run-20220609_163612-2jh9cxyp<code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>/tmp/tmpcbvaf5w6/wandb/offline-run-20220609_163612-2jh9cxyp/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#| hide\n",
    "#| slow\n",
    "path = untar_data(URLs.MNIST_TINY)\n",
    "items = get_image_files(path)\n",
    "tds = Datasets(items, [PILImageBW.create, [parent_label, Categorize()]], splits=GrandparentSplitter()(items))\n",
    "dls = tds.dataloaders(after_item=[ToTensor(), IntToFloatTensor()])\n",
    "\n",
    "os.environ['WANDB_MODE'] = 'dryrun' # run offline\n",
    "with tempfile.TemporaryDirectory() as wandb_local_dir:\n",
    "    wandb.init(anonymous='allow', dir=wandb_local_dir)\n",
    "    learn = vision_learner(dls, resnet18, loss_func=CrossEntropyLossFlat(), cbs=WandbCallback(log_model=False))\n",
    "    learn.fit(1)\n",
    "\n",
    "    # add more data from a new learner on same run\n",
    "    learn = vision_learner(dls, resnet18, loss_func=CrossEntropyLossFlat(), cbs=WandbCallback(log_model=False))\n",
    "    learn.fit(1, lr=slice(0.005))\n",
    "    \n",
    "    # save model\n",
    "    learn = cnn_learner(dls, resnet18, loss_func=CrossEntropyLossFlat(), cbs=WandbCallback(log_model=True))\n",
    "    learn.fit(1, lr=slice(0.005))\n",
    "    \n",
    "    # save model override name\n",
    "    learn = cnn_learner(dls, resnet18, path=wandb_local_dir,loss_func=CrossEntropyLossFlat(), cbs=[WandbCallback(log_model=True, model_name=\"good_name\"), SaveModelCallback(fname=\"bad_name\")])\n",
    "    learn.fit(1, lr=slice(0.005))\n",
    "    assert (Path(wandb_local_dir)/\"models/good_name.pth\").exists(), \"No model file found\"\n",
    "    \n",
    "    # finish writing files to temporary folder\n",
    "    wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "_all_ = ['wandb_process']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export -"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 00_torch_core.ipynb.\n",
      "Converted 01_layers.ipynb.\n",
      "Converted 01a_losses.ipynb.\n",
      "Converted 02_data.load.ipynb.\n",
      "Converted 03_data.core.ipynb.\n",
      "Converted 04_data.external.ipynb.\n",
      "Converted 05_data.transforms.ipynb.\n",
      "Converted 06_data.block.ipynb.\n",
      "Converted 07_vision.core.ipynb.\n",
      "Converted 08_vision.data.ipynb.\n",
      "Converted 09_vision.augment.ipynb.\n",
      "Converted 09b_vision.utils.ipynb.\n",
      "Converted 09c_vision.widgets.ipynb.\n",
      "Converted 10_tutorial.pets.ipynb.\n",
      "Converted 10b_tutorial.albumentations.ipynb.\n",
      "Converted 11_vision.models.xresnet.ipynb.\n",
      "Converted 12_optimizer.ipynb.\n",
      "Converted 13_callback.core.ipynb.\n",
      "Converted 13a_learner.ipynb.\n",
      "Converted 13b_metrics.ipynb.\n",
      "Converted 14_callback.schedule.ipynb.\n",
      "Converted 14a_callback.data.ipynb.\n",
      "Converted 15_callback.hook.ipynb.\n",
      "Converted 15a_vision.models.unet.ipynb.\n",
      "Converted 16_callback.progress.ipynb.\n",
      "Converted 17_callback.tracker.ipynb.\n",
      "Converted 18_callback.fp16.ipynb.\n",
      "Converted 18a_callback.training.ipynb.\n",
      "Converted 18b_callback.preds.ipynb.\n",
      "Converted 19_callback.mixup.ipynb.\n",
      "Converted 20_interpret.ipynb.\n",
      "Converted 20a_distributed.ipynb.\n",
      "Converted 20b_tutorial.distributed.ipynb.\n",
      "Converted 21_vision.learner.ipynb.\n",
      "Converted 22_tutorial.imagenette.ipynb.\n",
      "Converted 23_tutorial.vision.ipynb.\n",
      "Converted 24_tutorial.image_sequence.ipynb.\n",
      "Converted 24_tutorial.siamese.ipynb.\n",
      "Converted 24_vision.gan.ipynb.\n",
      "Converted 30_text.core.ipynb.\n",
      "Converted 31_text.data.ipynb.\n",
      "Converted 32_text.models.awdlstm.ipynb.\n",
      "Converted 33_text.models.core.ipynb.\n",
      "Converted 34_callback.rnn.ipynb.\n",
      "Converted 35_tutorial.wikitext.ipynb.\n",
      "Converted 37_text.learner.ipynb.\n",
      "Converted 38_tutorial.text.ipynb.\n",
      "Converted 39_tutorial.transformers.ipynb.\n",
      "Converted 40_tabular.core.ipynb.\n",
      "Converted 41_tabular.data.ipynb.\n",
      "Converted 42_tabular.model.ipynb.\n",
      "Converted 43_tabular.learner.ipynb.\n",
      "Converted 44_tutorial.tabular.ipynb.\n",
      "Converted 45_collab.ipynb.\n",
      "Converted 46_tutorial.collab.ipynb.\n",
      "Converted 50_tutorial.datablock.ipynb.\n",
      "Converted 60_medical.imaging.ipynb.\n",
      "Converted 61_tutorial.medical_imaging.ipynb.\n",
      "Converted 65_medical.text.ipynb.\n",
      "Converted 70_callback.wandb.ipynb.\n",
      "Converted 70a_callback.tensorboard.ipynb.\n",
      "Converted 70b_callback.neptune.ipynb.\n",
      "Converted 70c_callback.captum.ipynb.\n",
      "Converted 70d_callback.comet.ipynb.\n",
      "Converted 74_huggingface.ipynb.\n",
      "Converted 97_test_utils.ipynb.\n",
      "Converted 99_pytorch_doc.ipynb.\n",
      "Converted dev-setup.ipynb.\n",
      "Converted app_examples.ipynb.\n",
      "Converted camvid.ipynb.\n",
      "Converted distributed_app_examples.ipynb.\n",
      "Converted migrating_catalyst.ipynb.\n",
      "Converted migrating_ignite.ipynb.\n",
      "Converted migrating_lightning.ipynb.\n",
      "Converted migrating_pytorch.ipynb.\n",
      "Converted migrating_pytorch_verbose.ipynb.\n",
      "Converted ulmfit.ipynb.\n",
      "Converted index.ipynb.\n",
      "Converted quick_start.ipynb.\n",
      "Converted tutorial.ipynb.\n"
     ]
    }
   ],
   "source": [
    "#| hide\n",
    "from nbdev import *\n",
    "nbdev_export()"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "split_at_heading": true
  },
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
